import os
import time
import numpy as np
from PIL import Image
from flask import Flask, request, jsonify
from flask_cors import CORS
import tensorflow as tf
from model.model_head_CBAM import efficientnetv2_m as create_model

# ========== 环境与 GPU 设置 ==========
gpu_id = os.environ.get("GPU_ID", "unknown")

# 显存按需增长，避免一次性占满
gpus = tf.config.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# ========== 基本参数 ==========
input_size = (384, 384)
num_classes = 15
label_names = [
    'colorResult_grey', 'colorResult_white', 'colorResult_yellow',
    'shapeResult_ToothMarks', 'shapeResult_fat', 'shapeResult_normal',
    'shapeResult_thin', 'textureResult_dark', 'textureResult_normal',
    'textureResult_tender', 'textureResult_water', 'thicknessResult_Stripping',
    'thicknessResult_ecchymosis', 'thicknessResult_greasy', 'thicknessResult_thin'
]

# === 加载模型（绑定当前可见 GPU 的第一张）===
# 注意：/GPU:0 表示“当前进程可见的第 0 张 GPU”
with tf.device('/GPU:0'):
    print(f"[INFO] 模型绑定 TensorFlow 逻辑设备 /GPU:0")
    model = create_model(num_classes=num_classes)
    model.build((None, *input_size, 3))
    model.load_weights("./efficientnetv2_best_head_CBAM.h5")
    print("[INFO] 模型加载完成")

# ========== 图像预处理 ==========
def preprocess_image_from_bytes(file_stream):
    image = Image.open(file_stream).convert('RGB').resize(input_size, Image.LANCZOS)
    img_array = np.asarray(image, dtype=np.float32) / 255.0
    return np.expand_dims(img_array, axis=0)

# ========== Flask 应用 ==========
app = Flask(__name__)
CORS(app)

@app.route('/upload', methods=['POST'])
def upload_image():
    # 检验图片数据
    if 'image' not in request.files:
        return jsonify({'success': False, 'message': '未找到图像部分'}), 400

    file = request.files['image']

    if file.filename == '':
        return jsonify({'success': False, 'message': '未选择文件'}), 400

    try:
        img_for_format_check = Image.open(file.stream)
        if img_for_format_check.format not in ('JPEG', 'PNG'):
            return jsonify({'success': False, 'message': '仅支持 JPG 或 PNG 格式的图像'}), 400
        # 格式检查后把流指针重置，供后续读取
        file.stream.seek(0)
    except Exception as e:
        return jsonify({'success': False, 'message': f'图像处理失败：{str(e)}'}), 400

    # 模型预测
    try:
        img_tensor = preprocess_image_from_bytes(file.stream)
        start_time = time.time()
        preds = model.predict(img_tensor, verbose=0)[0]
        end_time = time.time()
        prediction = {name: float(f"{prob:.2f}") for name, prob in zip(label_names, preds)}
        device_used = model.trainable_variables[0].device  # 记录使用的设备信息

        return jsonify({
            'success': True,
            'gpu': gpu_id,
            'device': device_used,
            'message': '成功预测',
            'inference_time_sec': round(end_time - start_time, 4),
            'results': prediction
        }), 200

    except Exception as e:
        return jsonify({'success': False, 'message': f'预测失败：{str(e)}'}), 500


if __name__ == '__main__':
    app.run(debug=True)